"""
HSP (Huairou Symmetry Platform) 数据模型
包含点群、空间群、磁性空间群、自旋空间群的数据模型定义

使用 Beanie ODM 进行 MongoDB 异步数据库操作
"""

from typing import Annotated, Generic, Literal, TypeVar

from beanie import Document, Indexed
from pydantic import BaseModel, Field

T = TypeVar("T")

# ==================== 基础模型 ====================


class SymmetryOperation(BaseModel):
    """对称操作模型

    包含完整的对称操作信息，遵循 Bilbao Crystallographic Server 的标准格式
    """

    id: int = Field(..., description="操作编号")
    coordinate_triplet: str = Field(..., description="坐标三元组表示（如 x,y,z 或 -y,x,z）")
    matrix: list[list[float]] = Field(..., description="3x3 或 4x4 对称操作矩阵（O3 矩阵）")
    ita_symbol: str = Field(..., description="International Tables for Crystallography 符号")
    seitz: str = Field(..., description="Seitz 符号")

    # 扩展字段：用于矩阵表示页面
    det: int | None = Field(None, description="行列式（±1）")
    axis: list[float] | None = Field(None, description="旋转轴单位向量 [nx, ny, nz]")
    angle: float | None = Field(None, description="旋转角度（弧度）")
    su2: list[list[dict[str, float]]] | None = Field(None, description="SU(2) 矩阵表示（复数格式：{real, imag}）")
    irrep_matrices: dict[str, list[list[float | dict[str, float]]]] | None = Field(
        None, description="各不可约表示的矩阵（irrep_label -> 矩阵）"
    )


class IrreducibleRepresentation(BaseModel):
    """不可约表示模型"""

    ir_label: str = Field(..., description="不可约表示标签")
    dimension: int = Field(..., description="表示维度")
    characters: list[float | str] = Field(..., description="特征标向量（数值或符号表示，如 ω, i, √2）")
    matrix_refs: list[str] | None = Field(None, description="矩阵表示引用")
    time_reversal: bool = Field(False, description="是否包含时间反演")
    is_spinor: bool = Field(False, description="是否为旋量表示")
    source: str | None = Field(None, description="数据来源")
    validation_status: Literal["verified", "computed", "pending"] = Field("pending", description="验证状态")


class CharacterTableIrrep(BaseModel):
    """特征标表中的不可约表示行"""

    gamma: str = Field(..., description="不可约表示标签（Γ 标记）")
    mulliken: str = Field(..., description="Mulliken 符号")
    characters: list[str] = Field(..., description="特征标值列表（数值或符号表示，如 ω, i, √2）")


class CharacterTable(BaseModel):
    """特征标表模型"""

    classes: list[str] = Field(..., description="共轭类列表")
    irreps: list[CharacterTableIrrep] = Field(..., description="不可约表示列表（第一行可能是 Mult. 表头）")
    operations: list[str] | None = Field(None, description="Bilbao 简写符号列表（如 ['1', '2z', 'my', 'mx']）")
    note: str | None = Field(None, description="特征标表备注")


class SubgroupInfo(BaseModel):
    """子群信息（Schoenflies + 指数）"""

    schoenflies: str = Field(..., description="子群 Schoenflies 符号")
    index: int = Field(..., description="子群指数（父群阶 / 子群阶）")


# ==================== 点群模型 ====================


class PointGroupBase(BaseModel):
    """点群基础模型"""

    group_id: str = Field(..., description="群唯一标识符")
    schoenflies: str = Field(..., description="Schoenflies 符号")
    international: str | None = Field(None, description="国际符号（Hermann-Mauguin）")
    name: str = Field(..., description="群名称（英文）")
    order: int | str = Field(..., description="群的阶（具体值或参数化表达式，如 2 或 'n'）")
    is_crystallographic: bool = Field(..., description="是否为晶体学点群")


class CrystallographicPointGroup(PointGroupBase):
    crystal_system: str = Field(..., description="晶系")
    laue_class: str = Field(..., description="Laue 类")


class NonCrystallographicPointGroup(PointGroupBase):
    family: Literal["Cn", "Cnv", "Cnh", "Sn", "Dn", "Dnh", "Dnd", "I", "Ih", "T", "Td", "Th", "O", "Oh"] = Field(
        ..., description="点群族"
    )
    parity: Literal["even", "odd"] | None = Field(None, description="n 的奇偶性")
    parameter_description: str | None = Field(None, description="参数描述（如 n 的取值范围或条件）")


class PointGroupDetail(PointGroupBase):
    """点群详细信息模型（API 响应模型）

    包含所有点群类型的完整信息，特有字段设为可选
    注：晶体学特有字段来自 CrystallographicPointGroup，非晶体学特有字段来自 NonCrystallographicPointGroup
    """

    # 晶体学点群特有字段（可选，来自 CrystallographicPointGroup）
    crystal_system: str | None = Field(None, description="晶系")
    laue_class: str | None = Field(None, description="Laue 类")

    # 非晶体学点群特有字段（可选，来自 NonCrystallographicPointGroup）
    family: Literal["Cn", "Cnv", "Cnh", "Sn", "Dn", "Dnh", "Dnd", "I", "Ih", "T", "Td", "Th", "O", "Oh"] | None = Field(
        None, description="点群族"
    )
    parity: Literal["even", "odd"] | None = Field(None, description="n 的奇偶性")
    parameter_description: str | None = Field(None, description="参数描述（如 n 的取值范围或条件）")

    # 详细信息字段
    subgroups: list[SubgroupInfo] = Field(default_factory=list, description="子群信息列表（含指数）")
    generators: list[int] = Field(default_factory=list, description="生成元对应的 operation id 列表")
    operations: list[SymmetryOperation] = Field(default_factory=list, description="对称操作列表")
    character_table: CharacterTable = Field(..., description="特征标表数据")

    representations: list[IrreducibleRepresentation] = Field(default_factory=list, description="不可约表示列表")
    supergroups: list[str] = Field(default_factory=list, description="超群ID列表")
    description: str | None = Field(None, description="群描述")
    references: list[str] = Field(default_factory=list, description="参考文献")


class PointGroupDocument(Document, PointGroupDetail):
    """点群数据库模型（Beanie Document）

    继承 PointGroupDetail 的所有字段，只覆盖需要添加索引的字段
    """

    group_id: Annotated[str, Indexed(unique=True)]
    schoenflies: Annotated[str, Indexed()]
    is_crystallographic: Annotated[bool, Indexed()]

    class Settings:
        name = "point_groups"
        indexes = [
            "schoenflies",
            [("is_crystallographic", 1), ("order", -1)],
            "crystal_system",
        ]


class PointGroupTreeEdge(BaseModel):
    """点群关系图中的边（类型级极大子群覆盖关系）"""

    source: str = Field(..., description="起点 group_id（父群）")
    target: str = Field(..., description="终点 group_id（子群，Schoenflies 类型）")
    index: int = Field(..., description="指数 [G:H]，若可用则返回")


class PointGroupTree(BaseModel):
    """点群关系图响应模型（用于 /symmetry/point-groups-tree）"""

    nodes: list[CrystallographicPointGroup] = Field(..., description="节点（32 个晶体点群类型）")
    edges: list[PointGroupTreeEdge] = Field(..., description="覆盖关系边集合（传递约简后）")


# ==================== 空间群模型 ====================


class WyckoffPosition(BaseModel):
    """Wyckoff 位置模型"""

    wyckoff_letter: str = Field(..., description="Wyckoff 字母")
    multiplicity: int = Field(..., description="多重度")
    site_symmetry: str = Field(..., description="位置点对称性")
    coordinates: list[str] = Field(..., description="坐标表示（符号形式）")


class SpaceGroupBase(BaseModel):
    """空间群基础模型"""

    group_id: str = Field(..., description="群唯一标识符")
    number: int = Field(..., description="空间群编号 (1-230)")
    international_symbol: str = Field(..., description="国际符号（Hermann-Mauguin）")
    international_full: str | None = Field(None, description="完整国际符号")
    schoenflies: str | None = Field(None, description="Schoenflies 符号")
    hall_symbol: str | None = Field(None, description="Hall 符号")
    name: str = Field(..., description="空间群名称（英文）")
    crystal_system: str = Field(..., description="晶系")
    point_group: str = Field(..., description="对应点群")


class SpaceGroupDetail(SpaceGroupBase):
    """空间群详细信息模型（API 响应模型）"""

    setting: str | None = Field(None, description="设置（如 hexagonal, rhombohedral）")
    lattice_system: str = Field(..., description="晶格系统")
    bravais_lattice: str = Field(..., description="Bravais 晶格类型")
    generators: list[str] = Field(default_factory=list, description="生成元")
    wyckoff_positions: list[WyckoffPosition] = Field(default_factory=list, description="Wyckoff 位置")
    representations: list[IrreducibleRepresentation] = Field(default_factory=list, description="k点处的不可约表示")
    subgroups: list[str] = Field(default_factory=list, description="子群ID列表")
    supergroups: list[str] = Field(default_factory=list, description="超群ID列表")
    latex_symbol: str | None = Field(None, description="LaTeX 格式符号")
    description: str | None = Field(None, description="群描述")
    references: list[str] = Field(default_factory=list, description="参考文献")


class SpaceGroupDocument(Document, SpaceGroupDetail):
    """空间群数据库模型（Beanie Document）

    继承 SpaceGroupDetail 的所有字段，只覆盖需要添加索引的字段
    """

    # 覆盖需要添加索引的字段
    group_id: Annotated[str, Indexed(unique=True)]  # type: ignore
    number: Annotated[int, Indexed()]  # type: ignore
    crystal_system: Annotated[str, Indexed()]  # type: ignore
    point_group: Annotated[str, Indexed()]  # type: ignore

    class Settings:
        name = "space_groups"
        indexes = [
            "number",
            "international_symbol",
            [("crystal_system", 1), ("number", 1)],
            "point_group",
        ]


# ==================== 磁性空间群模型 ====================


class MagneticSpaceGroupBase(BaseModel):
    """磁性空间群基础模型"""

    group_id: str = Field(..., description="群唯一标识符")
    og_number: int = Field(..., description="OG编号 (1-1651)")
    og_symbol: str = Field(..., description="OG符号")
    bns_number: str | None = Field(None, description="BNS编号")
    bns_symbol: str | None = Field(None, description="BNS符号")
    name: str = Field(..., description="磁性空间群名称（英文）")
    og_type: Literal["Type-I", "Type-II", "Type-III", "Type-IV"] = Field(..., description="OG分类类型")
    parent_space_group: int = Field(..., description="母空间群编号")


class MagneticSpaceGroupDetail(MagneticSpaceGroupBase):
    """磁性空间群详细信息模型（API 响应模型）"""

    unitary_operations: list[SymmetryOperation] = Field(default_factory=list, description="幺正对称操作")
    anti_unitary_operations: list[SymmetryOperation] = Field(default_factory=list, description="反幺正对称操作")
    representations: list[IrreducibleRepresentation] = Field(default_factory=list, description="磁性表示")
    magnetic_point_group: str | None = Field(None, description="磁性点群")
    time_reversal_invariant: bool = Field(False, description="是否具有时间反演不变性")
    description: str | None = Field(None, description="群描述")
    references: list[str] = Field(default_factory=list, description="参考文献")


class MagneticSpaceGroupDocument(Document, MagneticSpaceGroupDetail):
    """磁性空间群数据库模型（Beanie Document）

    继承 MagneticSpaceGroupDetail 的所有字段，只覆盖需要添加索引的字段
    """

    # 覆盖需要添加索引的字段
    group_id: Annotated[str, Indexed(unique=True)]  # type: ignore
    og_number: Annotated[int, Indexed()]  # type: ignore
    og_type: Annotated[Literal["Type-I", "Type-II", "Type-III", "Type-IV"], Indexed()]  # type: ignore
    parent_space_group: Annotated[int, Indexed()]  # type: ignore

    class Settings:
        name = "magnetic_space_groups"
        indexes = [
            "og_number",
            "og_type",
            [("og_type", 1), ("og_number", 1)],
            "parent_space_group",
        ]


# ==================== 自旋空间群模型 ====================


class SpinSpaceGroupBase(BaseModel):
    """自旋空间群基础模型"""

    group_id: str = Field(..., description="群唯一标识符")
    ssg_number: int = Field(..., description="自旋空间群编号")
    international_symbol: str = Field(..., description="国际符号")
    name: str = Field(..., description="自旋空间群名称（英文）")
    parent_msg: int | None = Field(None, description="母磁性空间群OG编号")
    parent_space_group: int = Field(..., description="母空间群编号")


class SpinSpaceGroupDetail(SpinSpaceGroupBase):
    """自旋空间群详细信息模型（API 响应模型）"""

    spin_operations: list[SymmetryOperation] = Field(default_factory=list, description="自旋对称操作")
    representations: list[IrreducibleRepresentation] = Field(default_factory=list, description="自旋表示")
    spin_point_group: str | None = Field(None, description="自旋点群")
    description: str | None = Field(None, description="群描述")
    references: list[str] = Field(default_factory=list, description="参考文献")


class SpinSpaceGroupDocument(Document, SpinSpaceGroupDetail):
    """自旋空间群数据库模型（Beanie Document）

    继承 SpinSpaceGroupDetail 的所有字段，只覆盖需要添加索引的字段
    """

    # 覆盖需要添加索引的字段
    group_id: Annotated[str, Indexed(unique=True)]  # type: ignore
    ssg_number: Annotated[int, Indexed()]  # type: ignore
    parent_space_group: Annotated[int, Indexed()]  # type: ignore

    class Settings:
        name = "spin_space_groups"
        indexes = [
            "ssg_number",
            "parent_space_group",
            "parent_msg",
        ]


# ==================== 查询参数模型 ====================


class PointGroupQueryParams(BaseModel):
    """点群查询参数"""

    type: Literal["crystallographic", "non-crystallographic", "all"] = Field("all", description="点群类型")
    family: str | None = Field(None, description="点群族（用于非晶体点群）")
    n: int | None = Field(None, description="参数n的值")
    parity: Literal["even", "odd", "any"] | None = Field(None, description="n的奇偶性")
    crystal_system: str | None = Field(None, description="晶系（用于晶体点群）")


class SpaceGroupQueryParams(BaseModel):
    """空间群查询参数"""

    number: int | None = Field(None, ge=1, le=230, description="空间群编号")
    crystal_system: str | None = Field(None, description="晶系")
    point_group: str | None = Field(None, description="点群")
    setting: str | None = Field(None, description="设置")


class MagneticGroupQueryParams(BaseModel):
    """磁性空间群查询参数"""

    og_number: int | None = Field(None, ge=1, le=1651, description="OG编号")
    og_type: Literal["Type-I", "Type-II", "Type-III", "Type-IV"] | None = Field(None, description="OG类型")
    parent_space_group: int | None = Field(None, description="母空间群编号")


class SpinGroupQueryParams(BaseModel):
    """自旋空间群查询参数"""

    ssg_number: int | None = Field(None, description="自旋空间群编号")
    parent_msg: int | None = Field(None, description="母磁性空间群编号")
    parent_space_group: int | None = Field(None, description="母空间群编号")


# ==================== 响应模型 ====================


class PaginationParams(BaseModel):
    """分页参数"""

    page: int = Field(1, ge=1, description="页码")
    page_size: int = Field(20, ge=1, le=100, description="每页大小")


class PaginatedResponse(BaseModel, Generic[T]):
    """分页响应（泛型）"""

    total: int = Field(..., description="总数")
    data: list[T] = Field(..., description="数据列表")
